import pygame, math

def blitRotateZoomXY(surf, original_image, origin, pivot, angle, scale_x, scale_y):

    image_rect = original_image.get_rect(topleft = (origin[0] - pivot[0], origin[1] - pivot[1]))
    offset_center_to_pivot = pygame.math.Vector2(origin) - image_rect.center
    offset_center_to_pivot.x *= scale_x
    offset_center_to_pivot.y *= scale_y
    rotated_offset = offset_center_to_pivot.rotate(-angle)
    rotated_image_center = (origin[0] - rotated_offset.x, origin[1] - rotated_offset.y)
    scaled_image = pygame.transform.smoothscale(original_image, (image_rect.width * scale_x, image_rect.height * scale_y))
    rotozoom_image = pygame.transform.rotate(scaled_image, angle)
    rect = rotozoom_image.get_rect(center = rotated_image_center)

    surf.blit(rotozoom_image, rect)

    return rotozoom_image, rect

class AnimatedObject:
    def __init__(self, x, y, width, height, rotation=0.0, pivot_point:tuple[int,int]|None=None, id:str|None=None):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.rotation = rotation
        if pivot_point is None:
            self.pivot_point = (width // 2, height // 2)
        else:
            self.pivot_point = pivot_point
        self.surface = pygame.surface.Surface((width, height), pygame.SRCALPHA)
        self.keyframes = []
        self.id = id

    def add_keyframe(self, time, position:tuple[int,int]|None=None, scale:tuple[int,int]|None=None, rotation:float|None=None):
        if position == None:
            if len(self.keyframes) == 0:
                position = (self.x, self.y)
            else:
                position = self.keyframes[len(self.keyframes)-1]['position']
        if scale == None:
            if len(self.keyframes) == 0:
                scale = (self.width, self.height)
            else:
                scale = self.keyframes[len(self.keyframes)-1]['scale']
        if rotation == None:
            if len(self.keyframes) == 0:
                rotation = self.rotation
            else:
                rotation = self.keyframes[len(self.keyframes)-1]['rotation']
        self.keyframes.append({
            'time': time,
            'position': position,
            'scale': scale,
            'rotation': rotation
        })

    def update(self, current_time):
        for i in range(len(self.keyframes) - 1):
            start_time = self.keyframes[i]['time']
            end_time = self.keyframes[i + 1]['time']

            start_pos = self.keyframes[i]['position']
            end_pos = self.keyframes[i + 1]['position']

            start_sca = self.keyframes[i]['scale']
            end_sca = self.keyframes[i + 1]['scale']

            start_rot = self.keyframes[i]['rotation']
            end_rot = self.keyframes[i + 1]['rotation']

            if start_time <= current_time <= end_time:
                time_ratio = (current_time - start_time) / (end_time - start_time)
                self.x = start_pos[0] + (end_pos[0] - start_pos[0]) * time_ratio
                self.y = start_pos[1] + (end_pos[1] - start_pos[1]) * time_ratio
                self.width = start_sca[0] + (end_sca[0] - start_sca[0]) * time_ratio
                self.height = start_sca[1] + (end_sca[1] - start_sca[1]) * time_ratio
                self.rotation = start_rot + (end_rot - start_rot) * time_ratio
                break

class AnimationEngine:
    def __init__(self, geometry:tuple[int,int], frame_count:int, frame_rate:int, bg_color:tuple[int,int,int]):
        self.geometry = geometry
        self.frame_count = frame_count
        self.frame_rate = frame_rate
        self.bg_color = bg_color
        self.anim_objs = {}
        self.images = {}
        self.actions = []

    def save_buffer(self, dir_path, buffer):
        for number, frame in enumerate(buffer):
            pygame.image.save(frame, f'{dir_path}/{number}.png')

    def render_buffer(self):
        buffer = []
        for number in range(self.frame_count):
            frame = self.render_frame(number)
            buffer.append(frame)
        return buffer

    def create_frame(self):
        frame = pygame.Surface(self.geometry)
        frame.fill(self.bg_color)
        return frame

    def render_frame(self, number):
        frame = self.create_frame()
        self.execute_actions(number, False)
        for image in sorted(self.images_to_blit_to_frame, key=lambda d: d['layering']['layer']):
            if image['layering']['objs'] == 'below':
                frame.blit(image['image'], image['position'])
        for anim_obj_id in self.anim_objs:
            anim_obj = self.update_anim_obj(anim_obj_id, number)
            scale_x = anim_obj.width / anim_obj.surface.get_width()
            scale_y = anim_obj.height / anim_obj.surface.get_height()
            if anim_obj_id in self.modified_anim_obj_images:
                blitRotateZoomXY(frame, self.modified_anim_obj_images[anim_obj_id], (anim_obj.x, anim_obj.y), anim_obj.pivot_point, anim_obj.rotation, scale_x, scale_y)
            else:
                blitRotateZoomXY(frame, anim_obj.surface, (anim_obj.x, anim_obj.y), anim_obj.pivot_point, anim_obj.rotation, scale_x, scale_y)
        self.execute_actions(number, True)
        for image in sorted(self.images_to_blit_to_frame, key=lambda d: d['layering']['layer']):
            if image['layering']['objs'] == 'above':
                frame.blit(image['image'], image['position'])
        return frame

    def update_anim_objs(self, frame_number):
        for anim_obj_id in self.anim_objs:
            self.update_anim_obj(anim_obj_id, frame_number)

    def get_anim_obj_rects_and_images(self, frame_number, return_object_ids=False):
        rects = []
        images = []
        if return_object_ids:
            object_ids = []
        frame = self.create_frame()
        for anim_obj_id in self.anim_objs:
            anim_obj = self.update_anim_obj(anim_obj_id, frame_number)
            scale_x = anim_obj.width / anim_obj.surface.get_width()
            scale_y = anim_obj.height / anim_obj.surface.get_height()
            image, rect = blitRotateZoomXY(frame, anim_obj.surface, (anim_obj.x, anim_obj.y), anim_obj.pivot_point, anim_obj.rotation, scale_x, scale_y)
            rects.append(rect)
            images.append(image)
            if return_object_ids:
                object_ids.append(anim_obj_id)
        if return_object_ids:
            return images, rects, object_ids
        else:
            return images, rects

    def update_anim_obj(self, object_number, frame_number):
        anim_obj = self.anim_objs[object_number]
        anim_obj.update(frame_number)
        return anim_obj

    def render_and_save_frame(self, frame_number, file_name):
        self.update_anim_objs(frame_number)
        frame = self.render_frame(frame_number)
        pygame.image.save(frame, file_name)

    def render_and_save_animation(self, dir_path):
        for number in range(self.frame_count):
            frame = self.render_frame(number)
            pygame.image.save(frame, f'{dir_path}/{number}.png')

    def execute_actions(self, current_time, render_below):
        self.modified_anim_obj_images = {}
        self.images_to_blit_to_frame = []
        for action in self.actions:
            if action['time'][0] <= current_time <= action['time'][1]:
                if action['type'] == 'blit':
                    image = self.images[action['image']]
                    image = pygame.transform.scale(image, (image.get_width() * action['scale'], image.get_height() * action['scale']))
                    if action['object']:
                        anim_obj = self.anim_objs[action['object']]
                        if not action['object'] in self.modified_anim_obj_images:
                            self.modified_anim_obj_images[action['object']] = anim_obj.surface.copy()
                        self.modified_anim_obj_images[action['object']].blit(image, action['position'])
                    else:
                        self.images_to_blit_to_frame.append({
                            'image': image,
                            'position': action['position']
                        })
                elif action['type'] == 'draw':
                    if action['draw'] == 'rectangle':
                        surface = pygame.surface.Surface(action['geometry'], pygame.SRCALPHA)
                        pygame.draw.rect(surface,
                            action['color'],
                            (0, 0, *action['geometry']),
                            action['width'] if action['width'] else 0,
                            action['border'] if action['border'] and isinstance(action['border'], int) else -1,
                            action['border'][0] if action['border'] and isinstance(action['border'], list) else -1,
                            action['border'][1] if action['border'] and isinstance(action['border'], list) else -1,
                            action['border'][2] if action['border'] and isinstance(action['border'], list) else -1,
                            action['border'][3] if action['border'] and isinstance(action['border'], list) else -1
                        )
                    elif action['draw'] == 'circle':
                        surface = pygame.surface.Surface([action['radius'] * 2, action['radius'] * 2], pygame.SRCALPHA)
                        pygame.draw.circle(surface,
                            action['color'],
                            [
                                action['position'][0] + action['radius'],
                                action['position'][1] + action['radius']
                            ],
                            action['radius'],
                            action['width'] if action['width'] else 0
                        )
                    elif action['draw'] == 'ellipse':
                        surface = pygame.surface.Surface(action['geometry'], pygame.SRCALPHA)
                        pygame.draw.ellipse(surface,
                            action['color'],
                            (0, 0, *action['geometry']),
                            action['width'] if action['width'] else 0
                        )
                    elif action['draw'] == 'arc':
                        surface = pygame.surface.Surface(action['geometry'], pygame.SRCALPHA)
                        pygame.draw.arc(surface,
                            action['color'],
                            (0, 0, *action['geometry']),
                            math.radians(action['angle']['start']) if action['angle']['degrees'] else action['angle']['start'],
                            math.radians(action['angle']['stop']) if action['angle']['degrees'] else action['angle']['stop'],
                            action['width']
                        )
                    elif action['draw'] == 'line':
                        surface = pygame.surface.Surface(
                            (
                                action['end'][0] - action['start'][0],
                                action['end'][1] - action['start'][1]
                            ),
                            pygame.SRCALPHA
                        )
                        pygame.draw.line(surface,
                            action['color'],
                            action['start'],
                            action['end'],
                            action['width']
                        )
                        action['position'] = action['start']
                    elif action['draw'] == 'lines':
                        surface = pygame.surface.Surface((self.geometry), pygame.SRCALPHA)
                        rect = pygame.draw.lines(surface,
                            action['color'],
                            False,
                            action['points'],
                            action['width']
                        )
                        action['position'] = (0,0)
                    elif action['draw'] == 'polygon':
                        surface = pygame.surface.Surface((self.geometry), pygame.SRCALPHA)
                        rect = pygame.draw.polygon(surface,
                            action['color'],
                            action['points'],
                            action['width'] if action['width'] else 0
                        )
                        action['position'] = (0,0)
                    if action['object']:
                        anim_obj = self.anim_objs[action['object']]
                        if not action['object'] in self.modified_anim_obj_images:
                            self.modified_anim_obj_images[action['object']] = anim_obj.surface.copy()
                        self.modified_anim_obj_images[action['object']].blit(surface, action['position'])
                    else:
                        self.images_to_blit_to_frame.append({
                            'image': surface,
                            'position': action['position'],
                            'layering': action['layering']
                        })
